package com.datareport.task;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.datareport.domain.entity.StatisticResult;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.time.LocalDateTime;
import java.util.Properties;

/**
 * 统计计算任务
 */
@Slf4j
public class StatisticsJob {
    public static void main(String[] args) throws Exception {
        // 解析参数
        String ruleId = args.length > 0 ? args[0] : null;
        if (ruleId == null) {
            log.error("规则ID不能为空");
            return;
        }
        
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        
        // 1. 配置检查点
        env.enableCheckpointing(60000); // 每分钟做一次检查点
        env.getCheckpointConfig().setCheckpointTimeout(30000);
        
        // 2. 读取数据源
        DataStream<String> sourceStream = env.addSource(new FlinkKafkaConsumer<>(
            "extract_topic",
            new SimpleStringSchema(),
            getKafkaProperties()
        ));
        
        // 3. 数据处理转换
        DataStream<StatisticResult> resultStream = sourceStream
            .map(new StatisticMapper())
            .keyBy(record -> record.getRuleId())
            .window(TumblingProcessingTimeWindows.of(Time.seconds(10)))
            .aggregate(new StatisticAggregateFunction());
        
        // 4. 结果输出
        resultStream.addSink(new StatisticSink());
        
        env.execute("Statistics Computation Job - Rule: " + ruleId);
    }
    
    /**
     * 获取Kafka配置
     */
    private static Properties getKafkaProperties() {
        Properties properties = new Properties();
        properties.setProperty("bootstrap.servers", "localhost:9092");
        properties.setProperty("group.id", "statistics-consumer");
        return properties;
    }
    
    /**
     * 统计数据映射
     */
    public static class StatisticMapper implements MapFunction<String, StatisticRecord> {
        @Override
        public StatisticRecord map(String value) throws Exception {
            try {
                JSONObject json = JSON.parseObject(value);
                
                StatisticRecord record = new StatisticRecord();
                record.setRuleId(json.getString("ruleId"));
                record.setGroupKey(json.getString("groupKey"));
                record.setValue(json.getDouble("value"));
                record.setTimestamp(System.currentTimeMillis());
                
                return record;
            } catch (Exception e) {
                log.error("解析数据失败: {}", e.getMessage());
                throw e;
            }
        }
    }
    
    /**
     * 统计记录
     */
    public static class StatisticRecord {
        private String ruleId;
        private String groupKey;
        private Double value;
        private Long timestamp;
        
        public String getRuleId() {
            return ruleId;
        }
        
        public void setRuleId(String ruleId) {
            this.ruleId = ruleId;
        }
        
        public String getGroupKey() {
            return groupKey;
        }
        
        public void setGroupKey(String groupKey) {
            this.groupKey = groupKey;
        }
        
        public Double getValue() {
            return value;
        }
        
        public void setValue(Double value) {
            this.value = value;
        }
        
        public Long getTimestamp() {
            return timestamp;
        }
        
        public void setTimestamp(Long timestamp) {
            this.timestamp = timestamp;
        }
    }
    
    /**
     * 统计累加器
     */
    public static class StatisticAccumulator {
        private String ruleId;
        private String groupKey;
        private double sum = 0;
        private double min = Double.MAX_VALUE;
        private double max = Double.MIN_VALUE;
        private long count = 0;
        
        public void update(StatisticRecord record) {
            this.ruleId = record.getRuleId();
            this.groupKey = record.getGroupKey();
            this.sum += record.getValue();
            this.min = Math.min(this.min, record.getValue());
            this.max = Math.max(this.max, record.getValue());
            this.count++;
        }
        
        public StatisticAccumulator merge(StatisticAccumulator other) {
            this.sum += other.sum;
            this.min = Math.min(this.min, other.min);
            this.max = Math.max(this.max, other.max);
            this.count += other.count;
            return this;
        }
        
        public StatisticResult toResult() {
            StatisticResult result = new StatisticResult();
            result.setRuleId(this.ruleId);
            result.setGroupKey(this.groupKey);
            result.setSumValue(this.sum);
            result.setMinValue(this.min);
            result.setMaxValue(this.max);
            result.setAvgValue(this.count > 0 ? this.sum / this.count : 0);
            result.setCountValue(this.count);
            result.setStatTime(LocalDateTime.now());
            result.setCreateTime(LocalDateTime.now());
            return result;
        }
    }
    
    /**
     * 统计聚合函数
     */
    public static class StatisticAggregateFunction 
            implements AggregateFunction<StatisticRecord, StatisticAccumulator, StatisticResult> {
        
        @Override
        public StatisticAccumulator createAccumulator() {
            return new StatisticAccumulator();
        }
        
        @Override
        public StatisticAccumulator add(StatisticRecord value, StatisticAccumulator acc) {
            // 根据规则配置执行聚合计算
            acc.update(value);
            return acc;
        }
        
        @Override
        public StatisticResult getResult(StatisticAccumulator acc) {
            return acc.toResult();
        }
        
        @Override
        public StatisticAccumulator merge(StatisticAccumulator a, StatisticAccumulator b) {
            return a.merge(b);
        }
    }
    
    /**
     * 统计结果输出
     */
    public static class StatisticSink implements SinkFunction<StatisticResult> {
        @Override
        public void invoke(StatisticResult result, Context context) {
            try {
                log.info("统计结果: ruleId={}, groupKey={}, sum={}, avg={}, min={}, max={}, count={}",
                        result.getRuleId(), result.getGroupKey(), result.getSumValue(), result.getAvgValue(),
                        result.getMinValue(), result.getMaxValue(), result.getCountValue());
                
                // 保存结果到数据库
                saveResultToDatabase(result);
            } catch (Exception e) {
                log.error("保存统计结果失败", e);
            }
        }
        
        private void saveResultToDatabase(StatisticResult result) {
            String url = "jdbc:mysql://localhost:3306/datareport?useSSL=false&serverTimezone=UTC";
            String username = "root";
            String password = "password";
            
            String sql = "INSERT INTO statistic_result (result_id, rule_id, group_key, sum_value, avg_value, " +
                    "max_value, min_value, count_value, stat_time, create_time) " +
                    "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
            
            try (Connection conn = DriverManager.getConnection(url, username, password);
                 PreparedStatement stmt = conn.prepareStatement(sql)) {
                
                stmt.setString(1, java.util.UUID.randomUUID().toString().replace("-", ""));
                stmt.setString(2, result.getRuleId());
                stmt.setString(3, result.getGroupKey());
                stmt.setDouble(4, result.getSumValue());
                stmt.setDouble(5, result.getAvgValue());
                stmt.setDouble(6, result.getMaxValue());
                stmt.setDouble(7, result.getMinValue());
                stmt.setLong(8, result.getCountValue());
                stmt.setObject(9, result.getStatTime());
                stmt.setObject(10, result.getCreateTime());
                
                stmt.executeUpdate();
                
                log.info("统计结果已保存到数据库");
            } catch (Exception e) {
                log.error("保存统计结果到数据库失败", e);
            }
        }
    }
}